import random

# 定义班次和护士数量
num_shifts = 7
num_nurses = 10

# 定义班次要求和护士要求
shift_requirements = [1, 2, 2, 3, 2, 1, 1]
nurse_requirements = [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]

population_size = 100
mutation_rate = 0.01
generations = 1000

# 初始化种群
def init_population():
    population = []
    for i in range(population_size):
        chromosome = []
        for j in range(num_nurses):
            gene = [random.randint(0, 1) for _ in range(num_shifts)]
            chromosome.append(gene)
        population.append(chromosome)
    return population

# 计算适应度
def fitness(chromosome):
    score = 0
    for i in range(num_shifts):
        shift_count = 0
        for j in range(num_nurses):
            shift_count += chromosome[j][i]
        score += abs(shift_count - shift_requirements[i])
    for i in range(num_nurses):
        nurse_count = sum(chromosome[i])
        score += abs(nurse_count - nurse_requirements[i])
    return score

# 选择操作
def selection(population):
    fitnesses = [fitness(chromosome) for chromosome in population]
    total_fitness = sum(fitnesses)
    probabilities = [fitness / total_fitness for fitness in fitnesses]
    selected = []
    for i in range(population_size):
        selected.append(random.choices(population, probabilities)[0])
    return selected

# 交叉操作
def crossover(parent1, parent2):
    child1 = []
    child2 = []
    for i in range(num_nurses):
        if random.random() < 0.5:
            child1.append(parent1[i])
            child2.append(parent2[i])
        else:
            child1.append(parent2[i])
            child2.append(parent1[i])
    return child1, child2

# 变异操作
def mutation(chromosome):
    for i in range(num_nurses):
        for j in range(num_shifts):
            if random.random() < mutation_rate:
                chromosome[i][j] = 1 - chromosome[i][j]
    return chromosome

# 运行遗传算法
def run_ga():
    population = init_population()
    for generation in range(generations):
        population = selection(population)
        new_population = []
        for i in range(0, population_size, 2):
            parent1 = population[i]
            parent2 = population[i+1]
            child1, child2 = crossover(parent1, parent2)
            child1 = mutation(child1)
            child2 = mutation(child2)
            new_population.append(child1)
            new_population.append(child2)
        population = new_population
    best_chromosome = min(population, key=fitness)
    return best_chromosome

# 输出结果
best_chromosome = run_ga()
print("Best chromosome: ", best_chromosome)
print("Fitness: ", fitness(best_chromosome))